package com.pangyangyang.elasticsearch.plugin.jdbc;

import com.pangyangyang.elasticsearch.plugin.exception.MessageRiverException;
import com.pangyangyang.elasticsearch.plugin.jdbc.constant.JDBCHandlerConstants;
import com.pangyangyang.elasticsearch.plugin.jdbc.pool.JDBCConnectionPool;
import com.pangyangyang.elasticsearch.plugin.river.Build;
import com.pangyangyang.elasticsearch.plugin.river.MessageRiverSettings;
import com.pangyangyang.elasticsearch.plugin.util.Assert;
import org.elasticsearch.common.logging.ESLogger;
import org.elasticsearch.common.logging.ESLoggerFactory;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Created by yangwang on 2015/1/20.
 */
public class JDBCConnectionManager {

    private static ESLogger logger = ESLoggerFactory.getLogger(Build.LOGGER_NAME);

    private static Map<String, JDBCConnectionPool> connectionPools = new HashMap<>();
    private static Map<String, AtomicInteger> connectionRC = new HashMap<>();

    public static synchronized void registerConnectionPool(MessageRiverSettings settings){

        String url = settings.getString(JDBCHandlerConstants.SETTING_KEY_URL);
        if (!connectionPools.containsKey(url)){
            String username = settings.getString(JDBCHandlerConstants.SETTING_KEY_USERNAME);
            String password = settings.getString(JDBCHandlerConstants.SETTING_KEY_PASSWORD);
            JDBCConnectionPool pool = new JDBCConnectionPool(url, username, password);
            connectionPools.put(url, pool);
            connectionRC.put(url, new AtomicInteger(0));
        }
        connectionRC.get(url).incrementAndGet();
    }

    public static void unregisterConnectionPool(String mysqlUrl){
        AtomicInteger rc = connectionRC.get(mysqlUrl);
        //TODO 不完善，删除时有插入动作会出问题
        if(rc.decrementAndGet() <= 0){
            connectionPools.get(mysqlUrl).shutdown();
            connectionPools.remove(mysqlUrl);
            connectionRC.remove(mysqlUrl);
        }
    }

    public static Connection acquireConnection(String mysqlUrl){

        JDBCConnectionPool pool = connectionPools.get(mysqlUrl);
        Assert.assertNotNull(pool, "No connection pool found for " + mysqlUrl);

        try{
            return pool.getConnection();
        } catch (SQLException sqlException){
            throw new MessageRiverException("get connection failed.", sqlException);
        }
    }

    public static void releaseConnection(Connection connection){
        try {
            String url = connection.getClientInfo(JDBCHandlerConstants.SETTING_KEY_URL);
            JDBCConnectionPool pool = connectionPools.get(url);
            Assert.assertNotNull(pool, "Connection pool not found for " + url);
            pool.releaseConnection(connection);

        } catch (SQLException e) {
            logger.error("get jdbc connection url failed.", e);
        }
    }

    public static void closeStatement(Statement statement){
        if(statement != null){
            try{
                statement.close();
            } catch (Throwable t){
                //quiet
            }
        }

    }

    public static void closeResultSet(ResultSet rs){
        if(rs != null){
            try{
                rs.close();
            } catch (Throwable t){
                //quiet
            }
        }
    }
}
